import socket
import ssl
import threading

# 目标服务器配置
TARGET_HOST = "mirrors.aliyun.com"
TARGET_PORT = 443
INDEX_PATH_PREFIX = "/pypi/simple/"  # 仅索引页面需要的前缀

# 本地代理配置
PROXY_HOST = "0.0.0.0"
PROXY_PORT = 8081


def handle_client(client_socket):
    # 接收客户端请求
    request = client_socket.recv(4096)
    if not request:
        client_socket.close()
        return

    # 打印客户端原始请求
    print("\n[客户端原始请求]".center(50, '-'))
    print(request.decode('utf-8', errors='replace'))

    # 解析并修改请求（关键：区分索引请求和包文件请求）
    try:
        request_lines = request.decode('utf-8').split('\r\n')
        modified_lines = []
        path = None

        # 处理第一行（请求行）提取路径
        if request_lines[0].startswith(('GET', 'POST', 'HEAD')):
            method, original_path, version = request_lines[0].split(' ', 2)

            # 路径处理逻辑：
            # 1. 包文件请求（/packages/...）直接转发，不加前缀
            # 2. 索引页面请求（如/pillow/）添加INDEX_PATH_PREFIX前缀
            if original_path.startswith('/packages/'):
                path = "/pypi/" + original_path.lstrip('/')  # 索引路径加前缀
                # path = INDEX_PATH_PREFIX + original_path.lstrip('/')  # 索引路径加前缀
                # path = original_path  # 保持原始路径
            else:
                path = INDEX_PATH_PREFIX + original_path.lstrip('/')  # 索引路径加前缀

            modified_lines.append(f"{method} {path} {version}")
        else:
            modified_lines.append(request_lines[0])

        # 处理其他头部（修正Host头）
        for line in request_lines[1:]:
            if line.startswith('Host:'):
                modified_lines.append(f"Host: {TARGET_HOST}")  # 修正为目标域名
            else:
                modified_lines.append(line)

        # 重新组合请求
        modified_request = '\r\n'.join(modified_lines).encode('utf-8')
    except Exception as e:
        print(f"[解析请求错误] {e}")
        modified_request = request

    # 打印修改后发送到服务器的请求
    print("\n[修改后发送到服务器的请求]".center(50, '-'))
    print(modified_request.decode('utf-8', errors='replace'))

    # 连接目标服务器并转发请求
    target_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    context = ssl.create_default_context()
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE  # 忽略SSL验证

    try:
        # 建立SSL连接
        ssl_socket = context.wrap_socket(target_socket, server_hostname=TARGET_HOST)
        ssl_socket.connect((TARGET_HOST, TARGET_PORT))

        # 发送修改后的请求
        ssl_socket.send(modified_request)

        # 双向转发数据
        while True:
            # 接收服务器响应并转发给客户端
            server_response = ssl_socket.recv(4096)
            if not server_response:
                break
            print("\n[服务器响应]".center(50, '-'))
            print(server_response.decode('utf-8', errors='replace'))
            client_socket.send(server_response)

        # 接收客户端后续数据并转发给服务器
        while True:
            client_data = client_socket.recv(4096)
            if not client_data:
                break
            print("\n[客户端后续数据]".center(50, '-'))
            print(client_data.decode('utf-8', errors='replace'))
            ssl_socket.send(client_data)

    except Exception as e:
        print(f"\n[转发错误] {str(e)}")
    finally:
        ssl_socket.close()
        client_socket.close()
        print("\n[连接关闭]".center(50, '-'))


def main():
    proxy_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        proxy_socket.bind((PROXY_HOST, PROXY_PORT))
        proxy_socket.listen(5)
        print(f"代理服务器启动：{PROXY_HOST}:{PROXY_PORT}")
        print(f"转发目标：https://{TARGET_HOST}")
        print("按Ctrl+C停止")

        while True:
            client_sock, addr = proxy_socket.accept()
            print(f"\n[新连接] 来自 {addr}")
            threading.Thread(target=handle_client, args=(client_sock,)).start()

    except PermissionError:
        print(f"权限错误：请以管理员身份运行或更换端口（当前端口{PROXY_PORT}）")
    except OSError as e:
        print(f"端口错误：{e}（可能端口已被占用）")
    except KeyboardInterrupt:
        print("\n服务器停止中...")
    finally:
        proxy_socket.close()


if __name__ == "__main__":
    main()
